# Copyright 2025 林芳崽. All rights reserved.

from typing import Callable, Union
from typing import Optional

import torch.utils.checkpoint
from torch import nn
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.integrations import use_kernel_forward_from_hub
from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_layers import (
    GenericForQuestionAnswering,
    GenericForTokenClassification,
    GenericForSequenceClassification,
    GradientCheckpointingLayer,
)
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from transformers.models.qwen3.modeling_qwen3 import apply_rotary_pos_emb, eager_attention_forward
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
from transformers.utils.deprecation import deprecate_kwarg
from transformers.utils.generic import check_model_inputs

# 只能用相对目录，使用相对目录导出到模型自动加载类中才能自动识别加载
from .configuration_ai_lab import AiLabConfig


@use_kernel_forward_from_hub("RMSNorm")
class AiLabRMSNorm(nn.Module):
    """
    RMSNorm（Root Mean Square Normalization,平方均值（方差的近似））归一化层，等价于T5中使用的T5LayerNorm。
    与传统LayerNorm的区别：仅对输入的平方均值进行归一化，不含均值中心化步骤，计算更高效。

    RMSNorm 与 LayerNorm 的区别：
        传统 LayerNorm 公式：(x - mean) / sqrt(var + eps) * weight + bias（包含均值中心化）
        RMSNorm 公式：x / sqrt(var + eps) * weight（无均值中心化，仅用平方均值归一化）
        优势：减少了均值计算步骤，速度更快，在 Transformer 类模型中表现与 LayerNorm 接近。

    作用：
        在 Transformer 等深度学习模型中，RMSNorm 用于缓解深度网络训练中的 “内部协变量偏移” 问题，通过归一化隐藏层特征的尺度，使梯度更稳定，加速模型收敛。由于计算更高效，常被用于大模型（如 T5、LLaMA 等）中替代 LayerNorm。
    """

    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        """
        AiLabRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        # 可学习的缩放参数，维度与隐藏层大小一致，初始化为1（不改变归一化后的值）
        self.weight = nn.Parameter(torch.ones(hidden_size))
        # 数值稳定性参数，避免分母为0（通常设为1e-6）
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        前向传播：对输入张量进行RMS归一化计算
        Args:
            hidden_states: 输入张量，形状通常为 [batch_size, seq_len, hidden_size]
        Returns:
            归一化并缩放后的张量，形状与输入一致
        """
        # 保存输入的原始数据类型（如float16/bfloat16），用于后续恢复
        input_dtype = hidden_states.dtype
        # 将输入转换为float32进行计算，避免低精度下的数值不稳定问题
        hidden_states = hidden_states.to(torch.float32)

        # 1. 计算输入在最后一个维度（hidden_size）上的平方均值（方差的近似）
        #    pow(2)：对每个元素求平方
        #    mean(-1, keepdim=True)：在最后一个维度上求均值，保持维度不变（便于广播）
        variance = hidden_states.pow(2).mean(-1, keepdim=True)

        # 2. 计算归一化因子：1 / sqrt(variance + eps)，实现数值稳定的归一化
        #    rsqrt等价于1/sqrt(x)，计算更高效
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)

        # 3. 用可学习的权重进行缩放，并恢复原始数据类型
        #    权重self.weight允许模型动态调整归一化后各维度的重要性
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


class AiLabMLP(nn.Module):
    """
    基于SwiGLU激活函数的前馈网络（MLP）实现，用于Transformer类模型中的特征非线性变换。
    核心作用是对注意力层输出的特征进行高维映射、非线性加工和维度还原，增强模型对复杂模式的表达能力。
    """

    def __init__(self, config):
        super().__init__()
        self.config = config
        # 模型隐藏层维度（输入/输出特征维度），即每个token的特征向量维度
        self.hidden_size = config.hidden_size
        # MLP中间层维度，用于特征升维以增强非线性表达能力，通常为hidden_size的2-4倍（或更大）
        self.intermediate_size = config.intermediate_size

        # 门控投影层：将输入特征映射到intermediate_size维度，用于控制信息流动（类似"开关"）
        # 原理：通过学习哪些特征需要被增强，哪些需要被抑制
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)

        # 上投影层：将输入特征映射到intermediate_size维度，用于生成待激活的特征
        # 与gate_proj配合，形成门控机制的两个分支
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)

        # 下投影层：将中间层特征从intermediate_size维度还原回hidden_size维度
        # 确保MLP输出维度与输入维度一致，便于后续残差连接和层归一化
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)

        # 激活函数：从配置中获取指定的非线性激活函数（如SwiGLU、GELU等）
        # 非线性激活是MLP增强表达能力的核心，能拟合复杂的非线性关系
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        """
        MLP前向传播过程，采用SwiGLU门控机制实现特征变换：
        1. 门控分支：gate_proj(x)通过激活函数生成"门控信号"（0-1之间的权重）
        2. 特征分支：up_proj(x)生成待处理的特征
        3. 门控操作：将门控信号与特征分支按元素相乘，动态筛选重要特征
        4. 维度还原：通过down_proj将筛选后的特征映射回hidden_size维度

        原理优势：相比传统MLP（线性→激活→线性），SwiGLU通过门控机制动态控制信息流动，
        既能保留重要特征，又能抑制噪声，在大语言模型中表现出更优的特征表达能力。

        Args:
            x: 输入特征张量，形状为 [batch_size, seq_len, hidden_size]

        Returns:
            输出特征张量，形状为 [batch_size, seq_len, hidden_size]（与输入维度一致）
        """

        # 1. 门控分支：gate_proj(x)通过激活函数生成"门控信号"（0-1之间的权重）
        gate_weight = self.gate_proj(x)
        act_weight = self.act_fn(gate_weight)
        # 2. 特征分支：up_proj(x)生成待处理的特征
        up_features = self.up_proj(x)
        # 3. 门控操作：将门控信号与特征分支按元素相乘，动态筛选重要特征
        y = act_weight * up_features
        down_proj = self.down_proj(y)
        return down_proj


class AiLabAttention(nn.Module):
    r"""
    因果注意力机制：其核心作用是确保模型在生成序列时只能 “看到” 过去的信息，而无法访问未来的信息，从而符合语言生成的因果逻辑

    缩放点积：在原始注意力分数矩阵上叠加一个下三角掩码矩阵:

    :math:`\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} + \text{mask} \right) V\)

    其中 mask 即为因果掩码矩阵，确保未来信息的注意力权重被清零。`

    线性变换：
    :math:`\text{output} = \text{input} \times W^T + b`

    其中：W 是可学习的权重矩阵，形状为 (out_features, in_features)（注意是输出维度在前）。
    b 是可学习的偏置向量（可选），形状为 (out_features)。
    运算逻辑：输入特征与权重矩阵的转置进行矩阵相乘，再加上偏置项（若启用）。
    """

    def __init__(self, config: AiLabConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx # 当前注意力层的索引
        # 计算每个注意力头的维度：默认使用隐藏层维度除以注意力头数
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        # 计算键值对组的数量：用于分组注意力机制
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        # 缩放因子：sqrt(d_k)的倒数，用于缩放QK^T的结果
        self.scaling = self.head_dim ** -0.5
        # 注意力 dropout 概率
        self.attention_dropout = config.attention_dropout
        # 标记为因果注意力（仅关注过去信息）
        self.is_causal = True

        # # 注意力总维度：注意力头数 × 每个头的维度
        attention_dim = config.num_attention_heads * self.head_dim
        # 定义查询、键、值的线性投影层
        # Q投影：将输入隐藏层维度映射到注意力总维度
        self.q_proj = nn.Linear(config.hidden_size, attention_dim, bias=config.attention_bias)
        # K投影：将输入隐藏层维度映射到键值头总维度
        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim,
                                bias=config.attention_bias)
        # V投影：与K投影维度一致
        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim,
                                bias=config.attention_bias)
        # 输出投影：将注意力结果映射回原始隐藏层维度
        self.o_proj = nn.Linear(attention_dim, config.hidden_size, bias=config.attention_bias)

        # 对Q和K进行RMS归一化（仅在头维度上）
        self.q_norm = AiLabRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
        self.k_norm = AiLabRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # thus post q_norm does not need reshape

        # 滑动窗口大小：仅在当前层为滑动窗口注意力时生效
        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
            self,
            hidden_states: torch.Tensor,                             #  输入张量，形状通常为 [batch_size, seq_len, hidden_size]
            position_embeddings: tuple[torch.Tensor, torch.Tensor],  # 位置编码 (cos, sin)
            attention_mask: Optional[torch.Tensor],                  # 注意力掩码，用于屏蔽无效token
            past_key_values: Optional[Cache] = None,                 # 缓存的历史键值对，用于加速生成过程
            cache_position: Optional[torch.LongTensor] = None,       # 缓存位置索引
            **kwargs: Unpack[FlashAttentionKwargs],                  # 其他注意力相关参数
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        # 获取输入形状（排除最后一个特征维度）
        input_shape = hidden_states.shape[:-1]
        # 计算隐藏层形状，增加注意力头维度：[batch_size, seq_len, num_heads, head_dim]
        hidden_shape = (*input_shape, -1, self.head_dim)

        # QKV线性投影与形状转换
        # Q投影：先通过线性层，再重塑为 [batch_size, seq_len, num_heads, head_dim]，最后转置为 [batch_size, num_heads, seq_len, head_dim]
        query_states_view=self.q_proj(hidden_states).view(hidden_shape)
        query_states = self.q_norm(query_states_view).transpose(1, 2)

        # K投影：与Q类似，但头数为键值头数
        key_states_view=self.k_proj(hidden_states).view(hidden_shape)
        key_states = self.k_norm(key_states_view).transpose(1, 2)
        # V投影：与K处理流程相同
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        # 应用旋转位置编码（RoPE）
        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        # 处理历史缓存（用于生成式任务，避免重复计算）
        if past_key_values is not None:
            # 准备缓存更新所需参数
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            # 更新缓存，将当前键值对加入历史缓存
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)

        # 选择注意力计算实现（默认使用基础实现）
        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        # 注意力权重：原始注意力分数（query 与 key 的相似度）
        attn_output, attn_weights = attention_interface(
            self,
            query_states,  # 查询张量 [batch_size, num_heads, seq_len, head_dim]
            key_states,    # 键张量 [batch_size, num_key_value_heads, seq_len, head_dim]
            value_states,  # 值张量 [batch_size, num_key_value_heads, seq_len, head_dim]
            attention_mask,  # 注意力掩码
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            sliding_window=self.sliding_window,  # main diff with Llama
            **kwargs,
        )
        # 将注意力输出重塑为原始输入形状，并通过输出投影层
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()  # [batch_size, seq_len, attention_dim]
        attn_output = self.o_proj(attn_output)  # 映射回 [batch_size, seq_len, hidden_size]
        return attn_output, attn_weights  # 返回注意力输出和权重（可选）


class AiLabDecoderLayer(GradientCheckpointingLayer):
    """
    解码器层实现，继承自梯度检查点层以支持内存优化

    每个解码器层包含两个主要子模块：
    1. 自注意力机制（Self-Attention）
    2. 多层感知机（MLP）
    采用Pre-LayerNorm结构，在每个子模块前应用归一化，后接残差连接
    """
    def __init__(self, config: AiLabConfig, layer_idx: int):
        """
        初始化解码器层

        Args:
           config: 模型配置对象，包含隐藏层维度、归一化参数等
           layer_idx: 当前层的索引，用于区分不同层的配置
        """
        super().__init__()
        self.hidden_size = config.hidden_size

        # 初始化自注意力模块
        self.self_attn = AiLabAttention(config=config, layer_idx=layer_idx)

        # 初始化前馈网络模块
        self.mlp = AiLabMLP(config)

        # 输入归一化层（用于自注意力模块前）
        self.input_layernorm = AiLabRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # 注意力后归一化层（用于MLP模块前）
        self.post_attention_layernorm = AiLabRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # 当前层使用的注意力类型（如滑动窗口注意力或普通注意力）
        self.attention_type = config.layer_types[layer_idx]

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
            self,
            hidden_states: torch.Tensor,                     # 输入张量，形状为 [batch_size, seq_len, hidden_size]
            attention_mask: Optional[torch.Tensor] = None,   # 注意力掩码，形状通常为 [batch_size, 1, seq_len, seq_len]
            position_ids: Optional[torch.LongTensor] = None, # 位置ID张量，形状为 [batch_size, seq_len]，用于位置编码
            past_key_values: Optional[Cache] = None,
            use_cache: Optional[bool] = False,
            cache_position: Optional[torch.LongTensor] = None,
            position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC，
            **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        """
        前向传播函数，执行解码器层的核心计算

        Args:
            hidden_states: 输入张量，形状为 [batch_size, seq_len, hidden_size]
            attention_mask: 注意力掩码，形状通常为 [batch_size, 1, seq_len, seq_len]，
                            用于屏蔽填充token或未来信息
            position_ids: 位置ID张量，形状为 [batch_size, seq_len]，用于位置编码
            past_key_values: 缓存的历史键值对，用于生成式任务加速
            use_cache: 是否缓存当前层的键值对
            cache_position: 缓存位置索引，指示新token在缓存中的位置
            position_embeddings: 预计算的位置嵌入 (cos, sin)，用于旋转位置编码
            **kwargs: 其他传递给注意力模块的参数

        Returns:
            经过当前解码器层处理后的张量，形状保持 [batch_size, seq_len, hidden_size]
        """
        # 自注意力子模块计算
        residual = hidden_states # 保存残差连接的原始输入
        hidden_states = self.input_layernorm(hidden_states) #  # 应用输入归一化

        #  调用自注意力机制
        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = residual + hidden_states  # 应用残差连接

        # 前馈网络子模块计算
        residual = hidden_states # 保存残差连接的输入（注意力输出）
        hidden_states = self.post_attention_layernorm(hidden_states)  # 应用注意力后归一化
        hidden_states = self.mlp(hidden_states)   # 调用MLP模块
        hidden_states = residual + hidden_states  # 应用残差连接
        return hidden_states


class AiLabRotaryEmbedding(nn.Module):
    """
    旋转位置编码（Rotary Position Embedding, RoPE）实现类

    核心功能：为注意力机制的Q/K张量添加位置依赖，通过旋转操作编码位置信息，
    保持相对位置编码的特性，且不引入额外的可学习参数
    """

    inv_freq: torch.Tensor  # fix linting for `register_buffer`

    def __init__(self, config: AiLabConfig, device=None):
        """
        初始化旋转位置编码模块。RoPE（旋转位置编码）的核心数学思想是：通过对注意力机制的 Q/K 张量进行旋转操作，将位置信息编码到向量的相位中，同时严格保持 “相对位置编码不变” 的关键特性，让模型能精准捕捉序列中 token 的相对位置依赖。

        Args:
            config: 模型配置对象，需包含max_position_embeddings、rope_scaling等参数
            device: 设备类型（如cpu、cuda），用于初始化缓冲区
        """
        super().__init__()
        #  向后兼容处理：优先从rope_scaling字典中获取rope_type，兼容旧版"type"键
        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
        else:
            self.rope_type = "default"

        self.max_seq_len_cached = config.max_position_embeddings  # 缓存的最大序列长度
        self.original_max_seq_len = config.max_position_embeddings  # 原始配置的最大序列长度
        self.config = config  # 保存配置对象

        # 获取对应类型的RoPE初始化函数
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        # 调用初始化函数，获取逆频率和注意力缩放因子
        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)

        # 注册逆频率为缓冲区（不参与梯度更新，persistent=False表示不保存到模型文件）
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq  # 保存原始逆频率，用于后续可能的动态更新

    @torch.no_grad() # 禁用梯度计算，节省内存且加速计算
    @dynamic_rope_update  # 应用动态RoPE更新装饰器（高级功能支持）
    def forward(self, x, position_ids):
        """
        前向传播：生成对应位置的RoPE编码（cos和sin张量）

        Args:
            x: 输入张量（通常是Q或K），形状为 [batch_size, seq_len, hidden_size] 或类似
            position_ids: 位置ID张量，形状为 [batch_size, seq_len]，标记每个token的位置索引

        Returns:
            cos: 余弦编码张量，形状与freqs一致， dtype与输入x一致
            sin: 正弦编码张量，形状与freqs一致， dtype与输入x一致
        """
        # 扩展逆频率维度：适配batch维度，形状从 [head_dim//2] → [batch_size, head_dim//2, 1]
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
        # 扩展位置ID维度：形状从 [batch_size, seq_len] → [batch_size, 1, seq_len]
        position_ids_expanded = position_ids[:, None, :].float()
        # 设备类型适配：避免mps设备的autocast问题，默认使用cpu的autocast配置
        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"

        # 强制使用float32计算，保证RoPE编码的数值稳定性
        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
            # 计算频率：inv_freq_expanded @ position_ids_expanded → [batch_size, head_dim//2, seq_len]
            # 转置后形状为 [batch_size, seq_len, head_dim//2]
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)

            # 拼接频率张量：将freqs复制一份并拼接，形状变为 [batch_size, seq_len, head_dim]
            emb = torch.cat((freqs, freqs), dim=-1)

            # 计算余弦和正弦编码，并应用注意力缩放因子
            cos = emb.cos() * self.attention_scaling
            sin = emb.sin() * self.attention_scaling

        # 转换为与输入x一致的数据类型，避免类型不匹配
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


@auto_docstring
class AiLabPreTrainedModel(PreTrainedModel):
    config: AiLabConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["AiLabDecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True

    _can_compile_fullgraph = True
    _supports_attention_backend = True
    _can_record_outputs = {
        "hidden_states": AiLabDecoderLayer,
        "attentions": AiLabAttention,
    }


@auto_docstring
class AiLabModel(AiLabPreTrainedModel):
    """
    自定义Transformer解码器模型，基于标准Transformer架构实现

    核心结构：嵌入层 → 多层解码器层 → 输出归一化层
    支持因果注意力、滑动窗口注意力、动态缓存等功能，适用于文本生成等任务
    """
    def __init__(self, config: AiLabConfig):
        """
        初始化模型主体结构

        Args:
            config: 模型配置对象，包含词汇表大小、隐藏层维度、层数等超参数
        """
        super().__init__(config)

        self.padding_idx = config.pad_token_id  # padding token的ID（用于嵌入层屏蔽padding）
        self.vocab_size = config.vocab_size  # 词汇表大小

        # 词嵌入层：将token ID映射为隐藏层维度的向量
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) # 对padding token的嵌入进行特殊处理（通常置零）
        # 解码器层列表：堆叠多个AiLabDecoderLayer构成深层网络
        self.layers = nn.ModuleList(
            [AiLabDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        # 输出归一化层：对解码器堆叠后的输出进行归一化
        self.norm = AiLabRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # 旋转位置编码模块：为注意力机制提供位置依赖的编码
        self.rotary_emb = AiLabRotaryEmbedding(config=config)
        # 梯度检查点开关：用于节省训练时的内存（默认关闭）
        self.gradient_checkpointing = False
        # 标记是否包含滑动窗口注意力层
        self.has_sliding_layers = "sliding_attention" in self.config.layer_types

        # 初始化权重并执行最终处理（继承自PreTrainedModel的方法）
        self.post_init()

    @check_model_inputs
    @auto_docstring
    def forward(
            self,
            input_ids: Optional[torch.LongTensor] = None,
            attention_mask: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.LongTensor] = None,
            past_key_values: Optional[Cache] = None,
            inputs_embeds: Optional[torch.FloatTensor] = None,
            use_cache: Optional[bool] = None,
            cache_position: Optional[torch.LongTensor] = None,
            **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutputWithPast:
        """
        前向传播函数，执行模型的核心计算流程

        Args:
            input_ids: token ID张量，形状为 [batch_size, seq_len]，与inputs_embeds二选一
            attention_mask: 注意力掩码张量，用于屏蔽padding或无效位置
            position_ids: 位置ID张量，形状为 [batch_size, seq_len]，标记每个token的位置
            past_key_values: 历史键值对缓存，用于生成式任务加速推理
            inputs_embeds: 预计算的token嵌入张量，形状为 [batch_size, seq_len, hidden_size]，与input_ids二选一
            use_cache: 是否缓存当前层的键值对用于后续推理
            cache_position: 缓存位置索引，标记新token在缓存中的位置
            **kwargs: 传递给解码器层的其他参数

        Returns:
            BaseModelOutputWithPast: 模型输出对象，包含：
                - last_hidden_state: 模型最终输出，形状为 [batch_size, seq_len, hidden_size]
                - past_key_values: 缓存的键值对（仅当use_cache=True时返回）
        """
        # 校验输入：input_ids和inputs_embeds必须且只能指定一个
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        # 若未提供预计算嵌入，则通过嵌入层将input_ids转换为嵌入向量
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        # 若启用缓存且未初始化缓存，则创建动态缓存对象
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        # 若未指定缓存位置，则自动计算：从历史缓存长度开始到当前序列长度结束
        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )
        # 若未指定位置ID，则使用缓存位置索引作为位置ID（扩展为batch维度）
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        # 处理注意力掩码：生成不同注意力类型对应的因果掩码
        # 若attention_mask不是字典（即未提前准备），则自动创建掩码
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            # Prepare mask arguments
            mask_kwargs = {
                "config": self.config,
                "input_embeds": inputs_embeds,
                "attention_mask": attention_mask,
                "cache_position": cache_position,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }

            # # 创建标准因果掩码（用于普通注意力层）
            causal_mask_mapping = {
                "full_attention": create_causal_mask(**mask_kwargs),
            }
            # 若包含滑动窗口注意力层，则额外创建滑动窗口因果掩码
            if self.has_sliding_layers:
                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)

        # 初始化隐藏状态为输入嵌入
        hidden_states = inputs_embeds

        # 生成旋转位置编码（所有解码器层共享）
        position_embeddings = self.rotary_emb(hidden_states, position_ids)
        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                **kwargs,
            )

        # 对解码器堆叠后的输出进行归一化
        hidden_states = self.norm(hidden_states)

        # 返回模型输出，包含最终隐藏状态和缓存（若启用）
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
        )


class AiLabForCausalLM(AiLabPreTrainedModel, GenerationMixin):
    """
    基于AiLab模型的因果语言模型(Causal Language Model)实现

    因果语言模型是一种自回归模型，只使用过去的token来预测下一个token，
    适用于文本生成等任务。
    """
    # 定义需要绑定权重的键名（lm_head和embed_tokens的权重通常绑定以节省参数）
    _tied_weights_keys = ["lm_head.weight"]
    # 定义张量并行计划
    _tp_plan = {"lm_head": "colwise_rep"}
    # 定义流水线并行计划
    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}

    def __init__(self, config: AiLabConfig):
        super().__init__(config)
        # 初始化基础模型（Transformer解码器）
        self.model = AiLabModel(config)
        # 词汇表大小
        self.vocab_size = config.vocab_size
        # 语言模型头部：将隐藏状态映射到词汇表概率分布
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # 初始化权重并应用最终处理
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
            self,
            input_ids: Optional[torch.LongTensor] = None,
            attention_mask: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.LongTensor] = None,
            past_key_values: Optional[Cache] = None,
            inputs_embeds: Optional[torch.FloatTensor] = None,
            labels: Optional[torch.LongTensor] = None,
            use_cache: Optional[bool] = None,
            cache_position: Optional[torch.LongTensor] = None,
            logits_to_keep: Union[int, torch.Tensor] = 0,
            **kwargs: Unpack[TransformersKwargs],
    ) -> CausalLMOutputWithPast:
        r"""
        前向传播函数

        Args:
            input_ids: 输入token的ID序列，形状为[batch_size, sequence_length]
            attention_mask: 注意力掩码，标记有效的token位置
            position_ids: 位置ID，用于指示每个token的位置信息
            past_key_values: 缓存的键值对，用于加速生成过程
            inputs_embeds: 已经嵌入的输入向量，如果提供则忽略input_ids
            labels: 标签，用于计算损失函数
            use_cache: 是否使用缓存机制
            cache_position: 缓存位置信息
            logits_to_keep: 需要计算的logits数量，用于优化计算

        Example:

        ```python
        >>> from transformers import AutoTokenizer,Qwen2ForCausalLM # AiLabForCausalLM

        >>> model = AiLabForCausalLM.from_pretrained("meta-qwen2/Qwen2-2-7b-hf")
        >>> tokenizer = AutoTokenizer.from_pretrained("meta-qwen2/Qwen2-2-7b-hf")

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."

        Returns:
            CausalLMOutputWithPast: 包含损失、logits、缓存等信息的输出对象
        ```"""
        # 调用基础模型进行前向传播
        outputs: BaseModelOutputWithPast = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            cache_position=cache_position,
            **kwargs,
        )
        # 获取最后一层的隐藏状态
        hidden_states = outputs.last_hidden_state

        # 仅计算必要的logits，如果不计算损失则不将其提升为float类型，这是一个优化，避免不必要的计算
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep

        # 如果提供了标签，则计算损失
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

        # 返回因果语言模型的输出结果
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class AiLabForQuestionAnswering(GenericForQuestionAnswering, AiLabPreTrainedModel):
    base_model_prefix = "transformer"  # For BC, where `transformer` was used instead of `model`


class AiLabForSequenceClassification(GenericForSequenceClassification, AiLabPreTrainedModel):
    pass


class AiLabForTokenClassification(GenericForTokenClassification, AiLabPreTrainedModel):
    pass


__all__ = [
    "AiLabForCausalLM",
    "AiLabForQuestionAnswering",
    "AiLabPreTrainedModel",
    "AiLabModel",
    "AiLabForSequenceClassification",
    "AiLabForTokenClassification",
]
